<template>
  <div class="common-layout" v-loading="loading">
    <el-container>
      <div id="webgl">
        <Menu ref="menuRef" :joints="robotJoints" :controlJoint="controlJoint" :isChangePose="isChangePose" @editPose="asyncPose" @sliderInput="sliderInput" @setRobotCollision="setRobotCollision" @usePresetPoint="usePresetPoint" @setFollowLine="setFollowLine" />
      </div>
    </el-container>
  </div>
</template>

<script setup>
import gsap from 'gsap';
import Menu from './components/Menu/index.vue';
import * as THREE from 'three';
// Controller
import { OrbitControls } from 'three/addons/controls/OrbitControls.js';
// URDF Model Parser
import URDFLoader from 'urdf-loader';
import { PointerURDFDragControls } from './urdf-loader/URDFDragControls.js';
import { onMounted, ref } from 'vue';
import { MeshPhongMaterial } from 'three';

// Actual model
let robot = {};
// Drag model
let dragRobot = null;
// Whether the joint pose is being modified
let isChangePose = ref(false);
// Menu component instance
const menuRef = ref();
// Global loading state
const loading = ref(true);

// Create scene -> camera -> renderer -> add camera to scene -> renderer renders scene and camera -> add renderer to DOM
let scene = '';
let camera = '';
let renderer = '';
// Orbit controls
let controls = '';

// List of controllable joints for the robot arm
const robotJoints = ref([]);
// Key of the currently controlled joint
const controlJoint = ref('');

// Path to the URDF model of the robot arm
const urdf = './urdf/urdf.urdf';

// Line instances
let curveDataList = [];

// State of line visibility (enabled/disabled)
const curveType = ref(true);

// Listen for line visibility changes
const setFollowLine = type => {
  // Save line state
  curveType.value = type;

  // If line visibility is disabled, remove line instances
  if (!type && curveDataList.length !== 0) {
    // Remove all lines
    curveDataList.forEach(item => scene.remove(item));
    // Clear line list
    curveDataList = [];
  }
};

// Set angles for each drag joint
const sliderInput = (joint, value) => {
  // Set joint angle
  joint.setJointValue(value);

  // Update flag
  isChangePose.value = true;

  // Save joint list (filter out non-operable joints)
  robotJoints.value = Object.entries(dragRobot.joints).filter(item => !item[0].includes('joint_y'));
};

/**
 * Calculate the difference between two coordinates
 * @param {Array} arr1 Coordinate array 1
 * @param {Array} arr2 Coordinate array 2
 * @return {Number} Coordinate difference
 */
function getCoordinateDifference(arr1, arr2) {
  // Extract x, y, z values from the coordinates
  let arr1XYZ = arr1.slice(0, 6);
  let arr2XYZ = arr2.slice(0, 6);
  // Number of elements to compare
  let num = 6;

  // Sum of squared differences
  let sum = 0;
  for (let i = 0; i < num; i++) {
    sum += Math.pow(arr1XYZ[i] - arr2XYZ[i], 2);
  }
  // Calculate the Euclidean distance
  return Math.sqrt(sum);
}

// Throttle function
function throttle(func, delay) {
  let lastCall = 0;
  return function (...args) {
    const now = new Date().getTime();
    if (now - lastCall >= delay) {
      func.apply(this, args);
      lastCall = now;
    }
  };
}

/**
 * Generate a line following the robot arm's end effector
 * @param {Object} model The model node to follow
 * @param {String|Hex} color Line color, default is red
 * @param {Object} deviation Offset values for the line relative to the end effector
 * @return {Object} Returns methods to update the line
 */
function followArmLine(model, color = 0xff0000, deviation = { x: 0, y: 0, z: 0 }) {
  // If line visibility is disabled, skip rendering
  if (!curveType.value) return { curveObject: null, updateLineFn: () => {}, throttledUpdateLineFn: () => {} };

  // Create an invisible sphere to track world coordinates
  const sphere = new THREE.Mesh(new THREE.SphereGeometry(0, 32, 16), new THREE.MeshBasicMaterial({}));

  // Add sphere to the end effector
  model.add(sphere);

  // Apply custom offset to the sphere
  sphere.position.x += deviation.x || 0;
  sphere.position.y += deviation.y || 0;
  sphere.position.z += deviation.z || 0;

  // World position of the end effector
  let worldPosition1 = new THREE.Vector3();

  // Get world position of the sphere
  sphere.getWorldPosition(worldPosition1);

  // Create key points for the line (initialized with two identical points)
  const curve = new THREE.CatmullRomCurve3([worldPosition1, worldPosition1], false, 'centripetal', 0.5);
  // Generate points for the line
  const points = curve.getPoints(100);
  // Create geometry from the points
  const geometry = new THREE.BufferGeometry().setFromPoints(points);
  // Create line material
  const material = new THREE.LineBasicMaterial({ color });
  // Create THREE line
  const curveObject = new THREE.Line(geometry, material);
  // Save line instance
  curveDataList.push(curveObject);
  // Add line to the scene
  scene.add(curveObject);

  // Flag to indicate if the latest point should replace the second point
  let coverLastPoint = false;

  // Method to update the line in real-time
  function updateLineFn() {
    // Skip if line visibility is disabled
    if (!curveType.value) return;

    // World position of the end effector
    let worldPosition2 = new THREE.Vector3();

    // Update sphere position
    sphere.getWorldPosition(worldPosition2);

    // Get the last point in the curve
    const lastPoint = curve.points[curve.points.length - 1];

    // Skip update if the distance is too small
    if (worldPosition2.distanceTo(lastPoint) === 0) return;

    // If there are only two points and the last point hasn't been replaced
    if (curve.points.length === 2 && !coverLastPoint) {
      // Set flag to indicate replacement
      coverLastPoint = true;
      // Replace the second point with the new position
      curve.points[1] = worldPosition2;
    } else {
      // Add new point to the curve
      curve.points.push(worldPosition2);
    }

    // Generate updated points for the curve
    const updatedPoints = curve.getPoints(50);

    // Update geometry with new points
    geometry.setFromPoints(updatedPoints);

    // Notify Three.js to update the geometry
    curveObject.geometry.verticesNeedUpdate = true;
  }

  // Throttle the update method to avoid excessive calls
  const throttledUpdateLineFn = throttle(updateLineFn, 150);

  // Return methods to update the line
  return { throttledUpdateLineFn, updateLineFn };
}

// Synchronize poses
const asyncPose = () => {
  // Get actual robot joint list (filter out non-operable joints)
  const baseRobotJoints = Object.entries(robot.joints).filter(item => !item[0].includes('joint_y'));

  // Get actual robot joint poses
  const basePose = baseRobotJoints.map(item => item[1].angle);
  // Get drag robot joint poses
  const dragPose = Object.entries(dragRobot.joints)
    .filter(item => !item[0].includes('joint_y'))
    .map(item => item[1].angle);

  // Calculate the difference between the poses
  const difference = getCoordinateDifference(basePose, dragPose);

  // Save joint list (filter out non-operable joints)
  robotJoints.value = Object.entries(robot.joints).filter(item => !item[0].includes('joint_y'));

  // Get the end effector model node
  const armEnd = robotJoints.value[robotJoints.value.length - 1][1];

  // Create a follow line and get its update methods
  const { throttledUpdateLineFn, updateLineFn } = followArmLine(armEnd, 0xff4500, { z: 0.07 });

  // Apply all joints
  baseRobotJoints.forEach(item => {
    // Destructure data
    const [name, joint] = item;

    // Create GSAP configuration object
    let gsapConfig = {
      // Animation duration
      duration: difference,
      // Smooth animation
      ease: 'none',
      // Target value
      angle: dragRobot.joints[name].angle,
      // Callback when animation starts
      onStart: updateLineFn,
      // Callback during animation update
      onUpdate: () => {
        // Update robot arm pose
        joint.setJointValue(jointData.angle);
        // Update joint pose array (filter out non-operable joints)
        robotJoints.value = Object.entries(robot.joints).filter(item => !item[0].includes('joint_y'));

        // Throttle update for the follow line
        throttledUpdateLineFn();
      },
      // Callback when animation completes
      onComplete: () => {
        // Update flag
        isChangePose.value = false;

        // Directly update the follow line
        updateLineFn();
      },
    };

    // Create a motion object, this value will be modified in real-time
    let jointData = { angle: joint.angle };

    // Start animation
    gsap.to(jointData, gsapConfig);
  });
};

// Apply preset points
const usePresetPoint = targetPose => {
  // Get actual robot joint list (filter out non-operable joints)
  const baseRobotJoints = Object.entries(robot.joints).filter(item => !item[0].includes('joint_y'));

  // Get actual robot joint poses
  const basePose = baseRobotJoints.map(item => item[1].angle);

  // Get drag robot joint poses
  const dragPose = Object.entries(dragRobot.joints)
    .filter(item => !item[0].includes('joint_y'))
    .map(item => item[1].angle);

  // Get the end effector model node
  const armEnd = robotJoints.value[robotJoints.value.length - 1][1];

  // Create a follow line and get its update methods
  const { throttledUpdateLineFn, updateLineFn } = followArmLine(armEnd, 0x3498db, { z: 0.07 });

  /**
   * Helper method to animate the robot arm to a preset point
   * @param {Array} poseList Current robot arm pose
   * @param {Array} targetPose Target pose
   * @param {Robot} robot Robot instance
   * @param {Function} updateCallback Callback during animation update
   * @return {Promise} Returns a promise to track animation completion
   */
  function moveArmToPose(poseList, targetPose, robot, updateCallback) {
    return new Promise(resolve => {
      // Calculate the difference between the current and target poses
      const difference = getCoordinateDifference(poseList, targetPose);

      // Skip if there is no difference
      if (difference === 0) return;

      // Create a motion object, this value will be modified in real-time
      let jointData = {};

      // Get joint keys (filter out non-operable joints)
      const jointKeys = Object.keys(robot.joints).filter(item => !item.includes('joint_y'));

      // Create GSAP configuration object
      let gsapConfig = {
        // Animation duration
        duration: difference,
        // Smooth animation
        ease: 'none',
        // Callback when animation starts
        onStart: updateLineFn,
        // Callback during animation update
        onUpdate: () => {
          // Update robot arm pose
          robot.setJointValues(jointData);

          // Callback during animation update
          updateCallback && typeof updateCallback === 'function' && updateCallback(jointData);
        },
        // Callback when animation completes
        onComplete: () => {
          // Resolve the promise
          resolve();

          // Directly update the follow line
          updateLineFn();
        },
      };

      // Apply joint keys to avoid issues if the URDF model changes
      jointKeys.forEach((item, index) => {
        // Set motion object data
        jointData[item] = poseList[index];
        // Set GSAP configuration data
        gsapConfig[item] = targetPose[index];
      });

      // Start animation
      gsap.to(jointData, gsapConfig);
    });
  }

  // Animate the actual robot arm to the preset point
  const animation1 = moveArmToPose(basePose, targetPose, robot, () => {
    // Update joint pose array (filter out non-operable joints)
    robotJoints.value = Object.entries(robot.joints).filter(item => !item[0].includes('joint_y'));
    // Throttle update for the follow line
    throttledUpdateLineFn();
  });
  // Animate the drag robot arm to the preset point
  const animation2 = moveArmToPose(dragPose, targetPose, dragRobot);

  // Wait for all animations to complete, then reset the preset point selection
  Promise.all([animation1, animation2]).then(() => (menuRef.value.presetPointName = ''));
};

// Collision material
const collisionMaterial = new MeshPhongMaterial({
  transparent: true,
  opacity: 0.35,
  shininess: 2.5,
  premultipliedAlpha: true,
  color: 0xffbe38,
  polygonOffset: true,
  polygonOffsetFactor: -1,
  polygonOffsetUnits: -1,
});

// Default material list
let defaultMaterial = [];

// Show/Hide collision model
const setRobotCollision = isShow => {
  // Skip if robot is not loaded
  if (!robot) return;
  robot.traverse(item => {
    // Skip axes helper collision model
    if (item.name === 'axesHelper') return;

    // Show collision model
    if (isShow) {
      // Save original material
      defaultMaterial.push(item.material);
      // Set collision material
      item.material = collisionMaterial;
    } else {
      // Restore original material
      item.material = defaultMaterial.shift();
    }
  });
};

// Initialize base setup
function initBase() {
  // Create scene
  scene = new THREE.Scene();
  // Set scene position
  scene.position.set(0, 0, 0);
  // Create camera
  camera = new THREE.PerspectiveCamera(45, window.innerWidth / window.innerHeight, 0.1, 100);
  // Set camera position
  camera.position.set(1.8890530090942945, 2.8937984841724917, -1.0377233676862105);
  // Add camera to the scene
  scene.add(camera);

  // Create renderer with antialiasing and logarithmic depth buffer
  renderer = new THREE.WebGLRenderer({ antialias: true, logarithmicDepthBuffer: true });

  // Set renderer size
  renderer.setSize(window.innerWidth, window.innerHeight);

  // Set renderer background color
  renderer.setClearColor('#DCDCDC');
}

// Add lights to the scene
function addLight() {
  const positions = [
    { x: 15, y: 15, z: 15 },
    { x: -15, y: -15, z: -15 },
  ];
  positions.forEach(pos => {
    // Create directional light
    const light = new THREE.DirectionalLight('#fff', 1);
    // Set light position
    light.position.set(pos.x, pos.x, pos.z);
    // Add light to the scene
    scene.add(light);
  });
}

// Initialize orbit controls
function initOrbitControls() {
  // Create orbit controls
  controls = new OrbitControls(camera, renderer.domElement);
  // Enable damping for smoother controls
  controls.enableDamping = true;
}

// Render loop
function render() {
  // Update renderer
  renderer.render(scene, camera);
  // Update controls
  controls.update();

  // Request next frame
  requestAnimationFrame(render);
}

// Add grid helper to the scene
function addHelpLine() {
  // Create grid helper
  const gridHelper = new THREE.GridHelper(10, 50);
  gridHelper.position.y -= 1;
  scene.add(gridHelper);
}

/**
 * Add XYZ axes helper
 * @param {Object} model The model to add the helper to
 * @param {Number} thickness Thickness of the lines
 * @param {Number} height Height of the lines
 * @param {Object} deviation Offset values for the lines
 * @param {Array} helperOptions Configuration options for the helper
 */
function addAxesHelper(model, thickness, height, deviation, helperOptions) {
  // Define axes helper configurations
  let axesHelperArr = helperOptions || [
    {
      color: '#00FFB5',
      direction: 'x',
      cylinderDirection: 'z',
      goneDirection: 'y',
      angle: -Math.PI / 2,
    },
    {
      color: '#2F88FE',
      direction: 'y',
      cylinderDirection: 'y',
      goneDirection: 'y',
      angle: -Math.PI / 2,
    },
    {
      color: '#F86C55',
      direction: 'z',
      cylinderDirection: 'x',
      goneDirection: 'y',
      angle: Math.PI / 2,
    },
  ];

  // Loop through axes configurations
  for (let i = 0; i < axesHelperArr.length; i++) {
    // Destructure configuration
    const { color, direction, cylinderDirection, goneDirection, angle } = axesHelperArr[i];
    // Scale thickness
    const crude = thickness / 100;

    // Create cylinder geometry
    const geometryCylinder = new THREE.CylinderGeometry(crude, crude, height, 32);
    // Create cylinder material
    const materialCylinder = new THREE.MeshBasicMaterial({ color });
    // Create cylinder mesh
    const meshCylinder = new THREE.Mesh(geometryCylinder, materialCylinder);
    // Rotate cylinder
    meshCylinder.rotation[cylinderDirection] = angle;
    // Position cylinder
    meshCylinder.position[direction] = height / 2;

    // Create cone geometry
    const geometryCone = new THREE.ConeGeometry(crude * 2, 0.05, 32);
    // Create cone material
    const materialCone = new THREE.MeshBasicMaterial({ color });
    // Create cone mesh
    const meshCone = new THREE.Mesh(geometryCone, materialCone);
    // Position cone
    meshCone.position[goneDirection] = height / 2;

    // Add cone to cylinder
    meshCylinder.add(meshCone);

    // Apply deviation if provided
    if (deviation) {
      meshCylinder.position.x += deviation.x || 0;
      meshCylinder.position.y += deviation.y || 0;
      meshCylinder.position.z += deviation.z || 0;
    }

    // Set custom names for the helpers
    meshCone.name = 'axesHelper';
    meshCylinder.name = 'axesHelper';

    // Add cylinder and cone to the model
    model.add(meshCylinder);
  }
}

// Load robot model
function initRobot() {
  // Create URDF loader instance
  const loader = new URDFLoader();
  // Enable collision parsing
  loader.parseCollision = true;

  // Check if a joint is controllable
  const isJoint = j => j.isURDFJoint && j.jointType !== 'fixed';

  // Highlight material
  const highlightMaterial = new THREE.MeshPhongMaterial({
    shininess: 10,
    color: '#ffffff',
    emissive: '#ffffff',
    emissiveIntensity: 0.25,
  });

  // Highlight the currently hovered joint
  const highlightLinkGeometry = (m, revert) => {
    const traverse = item => {
      // Set or revert highlight color
      if (item.type === 'Mesh') {
        if (revert) {
          item.material = item.__origMaterial;
          delete item.__origMaterial;
        } else {
          item.__origMaterial = item.material;
          item.material = highlightMaterial;
        }
      }

      if (item === m || !isJoint(item)) {
        for (let i = 0; i < item.children.length; i++) {
          const child = item.children[i];
          if (!child.isURDFCollider) {
            traverse(item.children[i]);
          }
        }
      }
    };
    traverse(m);
  };

  // Load drag robot URDF model
  loader.load(urdf, urdfModel => {
    // Save URDF model
    dragRobot = urdfModel;

    // Create drag controls for the URDF model
    const dragControls = new PointerURDFDragControls(dragRobot, camera, renderer.domElement);

    // Disable orbit controls during drag
    dragControls.onDragStart = () => (controls.enabled = false);
    // Re-enable orbit controls after drag
    dragControls.onDragEnd = () => (controls.enabled = true);

    // Update joint angle during drag
    dragControls.updateJoint = (joint, angle) => sliderInput(joint, angle);

    // Highlight joint on hover
    dragControls.onHover = joint => {
      controlJoint.value = joint.name;
      highlightLinkGeometry(joint, false);
    };

    // Remove highlight on unhover
    dragControls.onUnhover = joint => {
      controlJoint.value = '';
      highlightLinkGeometry(joint, true);
    };

    // Rotate drag robot to align Z-axis upwards
    dragRobot.rotation.x = -Math.PI / 2;
    // Move model downwards
    dragRobot.position.y -= 1;

    // Save joint list (filter out non-operable joints)
    robotJoints.value = Object.entries(dragRobot.joints).filter(item => !item[0].includes('joint_y'));

    // Delay setting collision material
    setTimeout(() => {
      dragRobot.traverse(item => (item.material = collisionMaterial));
      // End loading
      loading.value = false;
    }, 500);

    // Add drag robot to the scene
    scene.add(dragRobot);
  });

  // Load initial URDF model
  loader.load(urdf, urdfModel => {
    // Save URDF model
    robot = urdfModel;

    // Rotate robot to align Z-axis upwards
    robot.rotation.x = -Math.PI / 2;
    // Move model downwards
    robot.position.y -= 1;

    // Get joint list
    let jointList = Object.keys(robot.joints).filter(item => !item.includes('joint_y'));
    // Last joint
    const lastJoint = jointList[jointList.length - 1];
    // First joint
    const firstJoint = jointList[0];

    // Add axes helper to the end effector
    addAxesHelper(robot.joints[lastJoint], 0.6, 0.15);

    // Add axes helper to the first joint
    addAxesHelper(robot.joints[firstJoint], 1, 0.3, { z: -0.1 });

    // Add robot to the scene
    scene.add(robot);
  });
}
initRobot();

// Initialize base setup
initBase();
// Add lights
addLight();
// Initialize orbit controls
initOrbitControls();
// Add grid helper
addHelpLine();

onMounted(() => {
  // Add renderer to the DOM
  document.getElementById('webgl').appendChild(renderer.domElement);
  render();

  // Handle window resize
  window.addEventListener('resize', () => {
    // Update camera aspect ratio
    camera.aspect = window.innerWidth / window.innerHeight;
    // Update camera projection matrix
    camera.updateProjectionMatrix();
    // Update renderer size
    renderer.setSize(window.innerWidth, window.innerHeight);
    // Set renderer pixel ratio
    renderer.setPixelRatio(window.innerWidth / window.innerHeight);
  });
});
</script>

<style>
.btn {
  position: fixed;
  bottom: 5%;
  left: 50%;
  transform: translateX(-50%);
}
</style>